package com.ai.module.websocket.handler;

import com.ai.module.chat.common.base.ThreadLocalContext;
import com.ai.module.chat.common.constant.RedisKeyStringConstant;
import com.ai.module.chat.common.constant.StringConstant;
import com.ai.module.chat.config.advice.exception.BusinessException;
import com.ai.module.chat.util.RedisTemplateUtils;
import com.ai.module.user.domain.vo.UserLoginInfo;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.Map;

/**
 * @author: heJinQin
 * @create: 2024-04-13 12:50
 * @description:
 **/
@Slf4j
@Component
public class WebSocketInterceptor implements HandshakeInterceptor {


    /**
     * 建立请求之前，可以用来做权限判断
     *
     * @param request    the current request
     * @param response   the current response
     * @param wsHandler  the target WebSocket handler
     * @param attributes the attributes from the HTTP handshake to associate with the WebSocket
     *                   session; the provided attributes are copied, the original map is not used.
     * @return
     * @throws Exception
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request,
                                   ServerHttpResponse response, org.springframework.web.socket.WebSocketHandler wsHandler,
                                   Map<String, Object> attributes) throws Exception {
        if (request instanceof ServletServerHttpRequest) {
            ServletServerHttpRequest servletServerHttpRequest = (ServletServerHttpRequest) request;
            // 模拟用户(通常利用JWT令牌解析用户信息)
            String headToken = servletServerHttpRequest.getServletRequest().getHeader(StringConstant.TOKEN);
            try {
                String userTokenKey = String.format(RedisKeyStringConstant.USER_INFO, headToken);
                UserLoginInfo userLoginInfo =  RedisTemplateUtils.get(userTokenKey);
                if (userLoginInfo ==null){
                    throw new BusinessException("登录已过期,请重新登录",401);
                }
                ThreadLocalContext.set(StringConstant.TOKEN, headToken);
                ThreadLocalContext.set(StringConstant.USER_INFO,userLoginInfo);
                // 设置当前这个session的属性，后续我们在发送消息时，可以通过 session.getAttributes().get("clientUserInfo")可以取出 clientUserInfo参数
                attributes.put(StringConstant.SOCKET_USER, userLoginInfo);
            } catch (Exception e) {
                log.error("webSocket 认证失败 ", e);
                return false;
            }
            return true;
        }
        return false;
    }

    /**
     * 建立请求之后
     *
     * @param request   the current request
     * @param response  the current response
     * @param wsHandler the target WebSocket handler
     * @param exception an exception raised during the handshake, or {@code null} if none
     */
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                               WebSocketHandler wsHandler, Exception exception) {

    }
}


